// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
import 'api.dart';
import 'content.dart';
import 'error.dart';

/// Configuration for a prebuilt voice.
///
/// This class allows specifying a voice by its name.
class PrebuiltVoiceConfig {
  // ignore: public_member_api_docs
  const PrebuiltVoiceConfig({this.voiceName});

  /// The voice name to use for speech synthesis.
  ///
  /// See https://cloud.google.com/text-to-speech/docs/chirp3-hd for names and
  /// sound demos.
  final String? voiceName;
  // ignore: public_member_api_docs
  Map<String, Object?> toJson() =>
      {if (voiceName case final voiceName?) 'voice_name': voiceName};
}

/// Configuration for the voice to be used in speech synthesis.
///
/// This class currently supports using a prebuilt voice configuration.
class VoiceConfig {
  // ignore: public_member_api_docs
  VoiceConfig({this.prebuiltVoiceConfig});

  // ignore: public_member_api_docs
  final PrebuiltVoiceConfig? prebuiltVoiceConfig;
  // ignore: public_member_api_docs
  Map<String, Object?> toJson() => {
        if (prebuiltVoiceConfig case final prebuiltVoiceConfig?)
          'prebuilt_voice_config': prebuiltVoiceConfig.toJson()
      };
}

/// Configures speech synthesis settings.
///
/// Allows specifying the desired voice for speech synthesis.
class SpeechConfig {
  /// Creates a [SpeechConfig] instance.
  ///
  /// [voiceName] See https://cloud.google.com/text-to-speech/docs/chirp3-hd
  /// for names and sound demos.
  SpeechConfig({String? voiceName})
      : voiceConfig = voiceName != null
            ? VoiceConfig(
                prebuiltVoiceConfig: PrebuiltVoiceConfig(voiceName: voiceName))
            : null;

  /// The voice config to use for speech synthesis.
  final VoiceConfig? voiceConfig;
  // ignore: public_member_api_docs
  Map<String, Object?> toJson() => {
        if (voiceConfig case final voiceConfig?)
          'voice_config': voiceConfig.toJson()
      };
}

/// The audio transcription configuration.
class AudioTranscriptionConfig {
  // ignore: public_member_api_docs
  Map<String, Object?> toJson() => {};
}

/// Configures live generation settings.
final class LiveGenerationConfig extends BaseGenerationConfig {
  // ignore: public_member_api_docs
  LiveGenerationConfig({
    this.speechConfig,
    this.inputAudioTranscription,
    this.outputAudioTranscription,
    super.responseModalities,
    super.maxOutputTokens,
    super.temperature,
    super.topP,
    super.topK,
    super.presencePenalty,
    super.frequencyPenalty,
  });

  /// The speech configuration.
  final SpeechConfig? speechConfig;

  /// The transcription of the input aligns with the input audio language.
  final AudioTranscriptionConfig? inputAudioTranscription;

  /// The transcription of the output aligns with the language code specified for
  /// the output audio.
  final AudioTranscriptionConfig? outputAudioTranscription;

  @override
  Map<String, Object?> toJson() => {
        ...super.toJson(),
        if (speechConfig case final speechConfig?)
          'speechConfig': speechConfig.toJson(),
      };
}

/// An abstract class representing a message received from a live server.
///
/// This class serves as a base for different types of server messages,
/// such as content updates, tool calls, and tool call cancellations.
/// Subclasses should implement specific message types.
sealed class LiveServerMessage {}

/// A message indicating that the live server setup is complete.
///
/// This message signals that the initial connection and setup process
/// with the live server has finished successfully.
class LiveServerSetupComplete implements LiveServerMessage {}

/// Audio transcription message.
class Transcription {
  // ignore: public_member_api_docs
  const Transcription({this.text, this.finished});

  /// Transcription text.
  final String? text;

  /// Whether this is the end of the transcription.
  final bool? finished;
}

/// Content generated by the model in a live stream.
class LiveServerContent implements LiveServerMessage {
  /// Creates a [LiveServerContent] instance.
  ///
  /// [modelTurn] (optional): The content generated by the model.
  /// [turnComplete] (optional): Indicates if the turn is complete.
  /// [interrupted] (optional): Indicates if the generation was interrupted.
  /// [inputTranscription] (optional): The input transcription.
  /// [outputTranscription] (optional): The output transcription.
  LiveServerContent(
      {this.modelTurn,
      this.turnComplete,
      this.interrupted,
      this.inputTranscription,
      this.outputTranscription});

  // TODO(cynthia): Add accessor for media content
  /// The content generated by the model.
  final Content? modelTurn;

  /// Whether the turn is complete. If true, indicates that the model is done
  /// generating.
  final bool? turnComplete;

  /// Whether generation was interrupted. If true, indicates that a
  /// client message has interrupted current model
  final bool? interrupted;

  /// The input transcription.
  ///
  /// The transcription is independent to the model turn which means it doesn't
  /// imply any ordering between transcription and model turn.
  final Transcription? inputTranscription;

  /// The output transcription.
  ///
  /// The transcription is independent to the model turn which means it doesn't
  /// imply any ordering between transcription and model turn.
  final Transcription? outputTranscription;
}

/// A tool call in a live stream.
///
/// A `Tool` is a piece of code that enables the system to interact with
/// external systems to perform an action, or set of actions, outside of
/// knowledge and scope of the model.
class LiveServerToolCall implements LiveServerMessage {
  /// Creates a [LiveServerToolCall] instance.
  ///
  /// [functionCalls] (optional): The list of function calls.
  LiveServerToolCall({this.functionCalls});

  /// The list of function calls to be executed.
  final List<FunctionCall>? functionCalls;
}

/// A tool call cancellation in a live stream.
///
/// Notification for the client that a previously issued `ToolCallMessage`
/// with the specified `id`s should have been not executed and should be
/// cancelled. If there were side-effects to those tool calls, clients may
/// attempt to undo the tool calls. This message occurs only in cases where the
/// clients interrupt server turns.
class LiveServerToolCallCancellation implements LiveServerMessage {
  /// Creates a [LiveServerToolCallCancellation] instance.
  ///
  /// [functionIds] (optional): The list of function IDs to cancel.
  LiveServerToolCallCancellation({this.functionIds});

  /// The list of [FunctionCall.id] to cancel.
  final List<String>? functionIds;
}

/// A single response chunk received during a live content generation.
///
/// It can contain generated content, function calls to be executed, or
/// instructions to cancel previous function calls, along with the status of the
/// ongoing generation.
class LiveServerResponse {
  // ignore: public_member_api_docs
  LiveServerResponse({required this.message});

  /// The server message generated by the live model.
  final LiveServerMessage message;
}

/// Represents realtime input from the client in a live stream.
class LiveClientRealtimeInput {
  /// Creates a [LiveClientRealtimeInput] instance.
  LiveClientRealtimeInput({
    @Deprecated('Use audio, video, or text instead') this.mediaChunks,
    this.audio,
    this.video,
    this.text,
  });

  /// Creates a [LiveClientRealtimeInput] with audio data.
  LiveClientRealtimeInput.audio(this.audio)
      // ignore: deprecated_member_use_from_same_package
      : mediaChunks = null,
        video = null,
        text = null;

  /// Creates a [LiveClientRealtimeInput] with video data.
  LiveClientRealtimeInput.video(this.video)
      // ignore: deprecated_member_use_from_same_package
      : mediaChunks = null,
        audio = null,
        text = null;

  /// Creates a [LiveClientRealtimeInput] with text data.
  LiveClientRealtimeInput.text(this.text)
      // ignore: deprecated_member_use_from_same_package
      : mediaChunks = null,
        audio = null,
        video = null;

  /// The list of media chunks.
  @Deprecated('Use audio, video, or text instead')
  final List<InlineDataPart>? mediaChunks;

  /// Audio data.
  final InlineDataPart? audio;

  /// Video data.
  final InlineDataPart? video;

  /// Text data.
  final String? text;

  // ignore: public_member_api_docs
  Map<String, dynamic> toJson() => {
        'realtime_input': {
          'media_chunks':
              // ignore: deprecated_member_use_from_same_package
              mediaChunks?.map((e) => e.toMediaChunkJson()).toList(),
          if (audio != null) 'audio': audio!.toMediaChunkJson(),
          if (video != null) 'video': video!.toMediaChunkJson(),
          if (text != null) 'text': text,
        },
      };
}

/// Represents content from the client in a live stream.
class LiveClientContent {
  /// Creates a [LiveClientContent] instance.
  ///
  /// [turns] (optional): The list of content turns from the client.
  /// [turnComplete] (optional): Indicates if the turn is complete.
  LiveClientContent({this.turns, this.turnComplete});

  /// The list of content turns from the client.
  final List<Content>? turns;

  /// Whether the turn is complete.
  ///
  /// If true, indicates that the server content generation should start with
  /// the currently accumulated prompt. Otherwise, the server will await
  /// additional messages before starting generation.
  final bool? turnComplete;

  // ignore: public_member_api_docs
  Map<String, dynamic> toJson() => {
        'client_content': {
          'turns': turns?.map((e) => e.toJson()).toList(),
          'turn_complete': turnComplete,
        }
      };
}

/// Represents a tool response from the client in a live stream.
class LiveClientToolResponse {
  /// Creates a [LiveClientToolResponse] instance.
  ///
  /// [functionResponses] (optional): The list of function responses.
  LiveClientToolResponse({this.functionResponses});

  /// The list of function responses.
  final List<FunctionResponse>? functionResponses;
  // ignore: public_member_api_docs
  Map<String, dynamic> toJson() => {
        'toolResponse': {
          'functionResponses': functionResponses
              ?.map((e) => {
                    'name': e.name,
                    'response': e.response,
                    if (e.id != null) 'id': e.id,
                  })
              .toList(),
        },
      };
}

/// Parses a JSON object received from the live server into a [LiveServerResponse].
///
/// This function handles different types of server messages, including:
/// - Error messages, which result in a [FirebaseAIException] being thrown.
/// - `serverContent` messages containing model-generated content.
/// - `toolCall` messages indicating function calls requested by the model.
/// - `toolCallCancellation` messages to cancel pending function calls.
/// - `setupComplete` messages signaling the completion of the server setup.
///
/// If the JSON object does not match any of the expected formats, an
/// [FirebaseAISdkException] is thrown.
///
/// Example:
/// ```dart
/// final jsonObject = {
///   'serverContent': {
///     'modelTurn': {
///       'parts': [
///         {'text': 'Hello, world!'}
///       ]
///     },
///     'turnComplete': true,
///   }
/// };
/// final message = parseServerMessage(jsonObject);
/// if (message is LiveServerContent) {
///   print('Received server content: ${message.modelTurn}');
/// }
/// ```
///
/// Throws:
/// - [FirebaseAIException]: If the JSON object contains an error message.
/// - [FirebaseAISdkException]: If the JSON object does not match any expected format.
///
/// Parameters:
/// - [jsonObject]: The JSON object received from the live server.
///
/// Returns:
/// - A [LiveServerResponse] object representing the parsed message.
LiveServerResponse parseServerResponse(Object jsonObject) {
  LiveServerMessage message = _parseServerMessage(jsonObject);
  return LiveServerResponse(message: message);
}

LiveServerMessage _parseServerMessage(Object jsonObject) {
  if (jsonObject case {'error': final Object error}) {
    throw parseError(error);
  }

  Map<String, dynamic> json = jsonObject as Map<String, dynamic>;

  if (json.containsKey('serverContent')) {
    final serverContentJson = json['serverContent'] as Map<String, dynamic>;
    Content? modelTurn;
    if (serverContentJson.containsKey('modelTurn')) {
      modelTurn = parseContent(serverContentJson['modelTurn']);
    }
    bool? turnComplete;
    if (serverContentJson.containsKey('turnComplete')) {
      turnComplete = serverContentJson['turnComplete'] as bool;
    }
    final interrupted = serverContentJson['interrupted'] as bool?;
    Transcription? _parseTranscription(String key) {
      if (serverContentJson.containsKey(key)) {
        final transcriptionJson =
            serverContentJson[key] as Map<String, dynamic>;
        return Transcription(
          text: transcriptionJson['text'] as String?,
          finished: transcriptionJson['finished'] as bool?,
        );
      }
      return null;
    }

    return LiveServerContent(
      modelTurn: modelTurn,
      turnComplete: turnComplete,
      interrupted: interrupted,
      inputTranscription: _parseTranscription('inputTranscription'),
      outputTranscription: _parseTranscription('outputTranscription'),
    );
  } else if (json.containsKey('toolCall')) {
    final toolContentJson = json['toolCall'] as Map<String, dynamic>;
    List<FunctionCall> functionCalls = [];
    if (toolContentJson.containsKey('functionCalls')) {
      final functionCallJsons =
          toolContentJson['functionCalls']! as List<dynamic>;
      for (final functionCallJson in functionCallJsons) {
        var functionCall =
            parsePart({'functionCall': functionCallJson}) as FunctionCall;
        functionCalls.add(functionCall);
      }
    }

    return LiveServerToolCall(functionCalls: functionCalls);
  } else if (json.containsKey('toolCallCancellation')) {
    final toolCancelData = json['toolCallCancellation'] as Map;
    final Map<String, List<String>> toolCancelJson = toolCancelData.map(
      (key, value) => MapEntry(
        key as String,
        (value as List).cast<String>(),
      ),
    );
    return LiveServerToolCallCancellation(functionIds: toolCancelJson['ids']);
  } else if (json.containsKey('setupComplete')) {
    return LiveServerSetupComplete();
  } else {
    throw unhandledFormat('LiveServerMessage', json);
  }
}
